/*
 * HapMatrix.cpp
 *
 *  Created on: Jan 20, 2012
 *      Author: pickrell
 */

#include "HapMatrix.h"
using namespace std;

HapMatrix::HapMatrix(string hapfile, string indfile, string snpfile){
	read_fastphase(hapfile);
	read_snpfile(snpfile);
	read_indfile(indfile);
	K = 100;
}

HapMatrix::HapMatrix(string hapfile, string snpfile, set<string> ind2keep){
	read_fastphase(hapfile, ind2keep);
	read_genmap(snpfile);
	K = 100;
}

void HapMatrix::read_fastphase(string hapfile){
	H.clear();
	string ext = hapfile.substr(hapfile.size()-3, 3);
	if (ext != ".gz"){
		std::cerr << hapfile << " is not gzipped (only .gz files accepted)\n";
		exit(1);
	}
	igzstream in(hapfile.c_str()); //only gzipped files
	vector<string> line;
	struct stat stFileInfo;
	int intStat;
	string st, buf;

	intStat = stat(hapfile.c_str(), &stFileInfo);
	if (intStat !=0){
		std::cerr<< "ERROR: cannot open file " << hapfile << "\n";
		exit(1);
	}
	for (int i = 0; i < 20; i++) getline(in, st);
	string p;
	int index;
	while(getline(in, st)){
		buf.clear();
		stringstream ss(st);
		line.clear();
		while (ss>> buf){
			line.push_back(buf);
		}
		string test = line[0];
		if (test.substr(0,1) == "H") {
			p = line[0];
			continue;
		}
		else if (test.substr(0,1) == "0" || test.substr(0,1) == "1"){
			vector<bool> tmp;
			for (vector<string>::iterator it = line.begin(); it != line.end(); it++){
				bool tmp2;
				if (*it == "0") tmp2 = false;
				else if (*it == "1") tmp2 = true;
				else {
					cerr << "ERROR: bad allele " << *it << "\n";
					exit(1);
				}
				tmp.push_back(tmp2);
			}
			H.push_back(tmp);
			index2ind.insert(make_pair(index, p));
			index++;
		}
	}
	nhap = H.size();
	nsnp = H.at(0).size();
}

void HapMatrix::read_fastphase(string hapfile, set<string> ind2keep){
	//cout << "here\n";
	H.clear();
	ind2index.clear();
	index2ind.clear();
	string ext = hapfile.substr(hapfile.size()-3, 3);
	if (ext != ".gz"){
		std::cerr << hapfile << " is not gzipped (only .gz files accepted)\n";
		exit(1);
	}
	igzstream in(hapfile.c_str()); //only gzipped files
	vector<string> line;
	struct stat stFileInfo;
	int intStat;
	string st, buf;

	intStat = stat(hapfile.c_str(), &stFileInfo);
	if (intStat !=0){
		std::cerr<< "ERROR: cannot open file " << hapfile << "\n";
		exit(1);
	}
	for (int i = 0; i < 20; i++) getline(in, st);
	string p;
	int index = 0;
	while(getline(in, st)){
		buf.clear();
		stringstream ss(st);
		line.clear();
		while (ss>> buf){
			line.push_back(buf);
		}
		//cout << line[0]<< "\n";
		if (line[0] == "#" and ind2keep.find(line[2]) != ind2keep.end()){
			p = line[2];
			ind2found[p] = true;
			//cout << p << "\n";
			string h1, h2;
			getline(in, h1);
			getline(in, h2);
			vector<string> line1;
			vector<string> line2;
			string buf1, buf2;
			stringstream ss1(h1);
			while (ss1 >> buf1){
				line1.push_back(buf1);
			}
			stringstream ss2(h2);
			while (ss2 >> buf2){
				line2.push_back(buf2);
			}
			vector<bool> tmp;
			vector<bool> tmp2;
			for (int i = 0; i < line1.size(); i++){
				bool tmp3;
				bool tmp4;
				if (line1[i] == "0") tmp3 = false;
				else if (line1[i] == "1") tmp3 = true;
				else {
					cerr << "ERROR: bad allele " << line1[i] << "\n";
					exit(1);
				}
				if (line2[i] == "0") tmp4 = false;
				else if (line2[i] == "1") tmp4 = true;
				else {
					cerr << "ERROR: bad allele " << line2[i] << "\n";
					exit(1);
				}
				tmp.push_back(tmp3);
				tmp2.push_back(tmp4);
			}
			H.push_back(tmp);
			H.push_back(tmp2);
			index2ind.insert(make_pair(index, p));
			index++;
			index2ind.insert(make_pair(index, p));
			ind2index.insert(make_pair(p, make_pair(index, index-1)));
			index++;
		}
	}
	nhap = H.size();
	nsnp = H.at(0).size();
}

void HapMatrix::print(){
	for (int i = 0; i < nhap; i++){
		string ind = index2ind[i];
		//string pop = ind2pop[ind];
		cout << ind << " ";
		for (int j = 0; j < nsnp; j++){
			cout << H[i][j] << " ";
		}
		cout << "\n";
	}
}
void HapMatrix::read_indfile(string indfile){
	ind2pop.clear();
	string ext = indfile.substr(indfile.size()-3, 3);
	if (ext != ".gz"){
		std::cerr << indfile << " is not gzipped (only .gz files accepted)\n";
		exit(1);
	}
	igzstream in(indfile.c_str()); //only gzipped files
	vector<string> line;
	struct stat stFileInfo;
	int intStat;
	string st, buf;

	intStat = stat(indfile.c_str(), &stFileInfo);
	if (intStat !=0){
		std::cerr<< "ERROR: cannot open file " << indfile << "\n";
		exit(1);
	}
	while(getline(in, st)){
		buf.clear();
		stringstream ss(st);
		line.clear();
		while (ss>> buf){
			line.push_back(buf);
		}
		allpops.insert(line[1]);
		ind2pop.insert(make_pair(line[0], line[1]));
	}
}

void HapMatrix::read_snpfile(string snpfile){
	index2snp.clear();
	snp2info.clear();
	string ext = snpfile.substr(snpfile.size()-3, 3);
	if (ext != ".gz"){
		std::cerr << snpfile << " is not gzipped (only .gz files accepted)\n";
		exit(1);
	}
	igzstream in(snpfile.c_str()); //only gzipped files
	vector<string> line;
	struct stat stFileInfo;
	int intStat;
	string st, buf;

	intStat = stat(snpfile.c_str(), &stFileInfo);
	if (intStat !=0){
		std::cerr<< "ERROR: cannot open file " << snpfile << "\n";
		exit(1);
	}
	getline(in, st);
	int index = 0;
	while(getline(in, st)){
		buf.clear();
		stringstream ss(st);
		line.clear();
		while (ss>> buf){
			line.push_back(buf);
		}
		index2snp.insert(make_pair(index, line[0]));
		snp2index.insert(make_pair(line[0], index));
		index++;
		vector<string> tmp;
		for (int i = 1; i < line.size(); i++){
			tmp.push_back(line[i]);
		}
		snp2info.insert(make_pair(line[0], tmp));
	}
}

void HapMatrix::read_genmap(string snpfile){
	index2snp.clear();
	snp2info.clear();
	genmap.clear();
	string ext = snpfile.substr(snpfile.size()-3, 3);
	if (ext != ".gz"){
		std::cerr << snpfile << " is not gzipped (only .gz files accepted)\n";
		exit(1);
	}
	igzstream in(snpfile.c_str()); //only gzipped files
	vector<string> line;
	struct stat stFileInfo;
	int intStat;
	string st, buf;

	intStat = stat(snpfile.c_str(), &stFileInfo);
	if (intStat !=0){
		std::cerr<< "ERROR: cannot open file " << snpfile << "\n";
		exit(1);
	}
	int index = 0;
	while(getline(in, st)){
		buf.clear();
		stringstream ss(st);
		line.clear();
		while (ss>> buf){
			line.push_back(buf);
		}
		index2snp.insert(make_pair(index, line[0]));
		snp2index.insert(make_pair(line[0], index));
		index++;
		vector<string> tmp;
		for (int i = 1; i < line.size(); i++){
			tmp.push_back(line[i]);
		}
		snp2info.insert(make_pair(line[0], tmp));
		genmap.push_back(atof(line[2].c_str()));
	}
	if (genmap.size() != nsnp){
		cerr << "ERROR: Number of SNPs in genetic map: "<< genmap.size() << ". In file: "<< nsnp << "\n";
		exit(1);
	}
}
map<string, map<string, vector<int> > > HapMatrix::get_pairwise_counts(set<string> pops){
	map<string, map<string, vector<int> > > toreturn;
	for (int i = 0; i < nsnp-1; i++){
		string snp1 = index2snp[i];
		int maxj = i+K;
		map<string, vector<int> > tmp;
		if (maxj > nsnp) maxj = nsnp;
		for (int j = i+1; j < maxj ; j++){
			string snp2 = index2snp[j];
			vector<int> tmp2;
			tmp2.push_back(0);tmp2.push_back(0);tmp2.push_back(0);tmp2.push_back(0);
			tmp.insert(make_pair(snp2, tmp2));
		}
		toreturn.insert(make_pair(snp1, tmp));
	}
	for (int i = 0; i < nhap; i++){

		string ind = index2ind[i];
		string pop = ind2pop[ind];
		if (pops.find(pop) == pops.end()) continue;
		cout << i << " "<< nhap << "\n";
		cout << ind << " "<< pop << "\n";
		for (int j = 0; j < nsnp-1; j++){
			string snp1 = index2snp[j];
			int maxk = j+K;
			if (maxk > nsnp) maxk = nsnp;
			for (int k = j+1; k < maxk ; k++){
				//cout << i << " "<< j << " "<< k << "\n";
				string snp2 = index2snp[k];
				//cout << snp1 << " "<< snp2 << "\n";
				bool a1 = H[i][j];
				bool a2 = H[i][k];
				//cout << a1 << " "<< a2 << "\n";
				int toadd;
				if (a1 == false and a2 == false) toadd = 0;
				else if (a1 == false and a2 == true) toadd = 1;
				else if (a1 == true and a2 == false) toadd = 2;
				else if (a1 == true and a2 == true) toadd = 3;

				//cout<< "here\n"; cout.flush();
				map< string, map<string, vector<int> > >::iterator it = toreturn.find(snp1);
				map<string, vector<int> >::iterator it2 = it->second.find(snp2);
				it2->second.at(toadd) = it2->second.at(toadd)+1;
				//cout << "not here\n"; cout.flush();
			}
		}
	}
	return toreturn;
}


map<string, map<string, vector<int> > > HapMatrix::get_pairwise_counts(set<string> pops, set<pair<string, string> > pairs){
	map<string, map<string, vector<int> > > toreturn;
	for (set<pair<string, string> >::iterator it = pairs.begin(); it != pairs.end(); it++){
		string snp1 = it->first;
		string snp2 = it->second;
		if (toreturn.find(snp1) == toreturn.end()){
			map<string, vector<int> > tmp;
			toreturn.insert(make_pair(snp1, tmp));
		}
		vector<int> tmp2;
		tmp2.push_back(0);tmp2.push_back(0);tmp2.push_back(0);tmp2.push_back(0);
		toreturn[snp1].insert(make_pair(snp2, tmp2));
	}


	for (int i = 0; i < nhap; i++){

		string ind = index2ind[i];
		string pop = ind2pop[ind];
		if (pops.find(pop) == pops.end()) continue;
		cout << i << " "<< nhap << "\n";
		cout << ind << " "<< pop << "\n";
		for (set<pair<string, string> >::iterator it = pairs.begin(); it != pairs.end(); it++){
			string snp1 = it->first;
			string snp2 = it->second;
			int i1 = snp2index[snp1];
			int i2 = snp2index[snp2];
			bool a1 = H[i][i1];
			bool a2 = H[i][i2];
			int toadd;
			if (a1 == false and a2 == false) toadd = 0;
			else if (a1 == false and a2 == true) toadd = 1;
			else if (a1 == true and a2 == false) toadd = 2;
			else if (a1 == true and a2 == true) toadd = 3;
			//cout<< "here\n"; cout.flush();
			map< string, map<string, vector<int> > >::iterator it = toreturn.find(snp1);
			map<string, vector<int> >::iterator it2 = it->second.find(snp2);
			it2->second.at(toadd) = it2->second.at(toadd)+1;
		}
	}
	return toreturn;
}
/*

map<string, map<string, map<string, vector<int> > > > HapMatrix::gt_pairs(set<string> pops, set<pair<string, string> > pairs){
	map<string, map<string, map<string, vector<int> > > > toreturn;
	for (set<pair<string, string> >::iterator it = pairs.begin(); it != pairs.end(); it++){
		string snp1 = it->first;
		string snp2 = it->second;
		if (toreturn.find(snp1) == toreturn.end()){
			map<string, map<string, vector<int> > > tmp;
			toreturn.insert(make_pair(snp1, tmp));
		}

		for(set<string>::iterator it2 = pops.begin(); it2 != pops.end(); it2++){
			map<string, vector<int> > tmp;
		}


		vector<int> tmp2;
		tmp2.push_back(0);tmp2.push_back(0);tmp2.push_back(0);tmp2.push_back(0);
		toreturn[snp1].insert(make_pair(snp2, tmp2));
	}


	for (int i = 0; i < nhap; i++){

		string ind = index2ind[i];
		string pop = ind2pop[ind];
		if (pops.find(pop) == pops.end()) continue;
		cout << i << " "<< nhap << "\n";
		cout << ind << " "<< pop << "\n";
		for (set<pair<string, string> >::iterator it = pairs.begin(); it != pairs.end(); it++){
			string snp1 = it->first;
			string snp2 = it->second;
			int i1 = snp2index[snp1];
			int i2 = snp2index[snp2];
			bool a1 = H[i][i1];
			bool a2 = H[i][i2];
			int toadd;
			if (a1 == false and a2 == false) toadd = 0;
			else if (a1 == false and a2 == true) toadd = 1;
			else if (a1 == true and a2 == false) toadd = 2;
			else if (a1 == true and a2 == true) toadd = 3;
			//cout<< "here\n"; cout.flush();
			map< string, map<string, vector<int> > >::iterator it = toreturn.find(snp1);
			map<string, vector<int> >::iterator it2 = it->second.find(snp2);
			it2->second.at(toadd) = it2->second.at(toadd)+1;
		}
	}
	return toreturn;
}

*/

map<string, pair<int, int> > HapMatrix::gt_snp(string snp1, string snp2, bool a1, bool a2){
	map<string, pair<int, int> > toreturn;
	for (int i = 0; i < nhap; i++){
		string ind = index2ind[i];
		string pop = ind2pop[ind];
		if (toreturn.find(pop) == toreturn.end()){
			pair<int, int> tmp = make_pair(0, 0);
			toreturn.insert(make_pair(pop, tmp));
		}
		int i1 = snp2index[snp1];
		int i2 = snp2index[snp2];
		bool ta1 = H[i][i1];
		bool ta2 = H[i][i2];
		if (ta1 == a1 and ta2 == a2) toreturn[pop].first ++;
		toreturn[pop].second ++;
	}
	return toreturn;
}

int HapMatrix::set_weights(set<string> mixinds, set<string> sourceinds, set<string> mix2, float * f2_uncor, float * hzy_mix, float * hzy_source, float * count_mix, float * count_source){
	snpinclude.clear();
	snpweight.clear();
	snpfreq.clear();

	int totalsnps = 0;
	for (int i = 0; i < nsnp; i++){
		int c1 = 0;
		int c2 = 0;
		int t1 = 0;
		int t2 = 0;
		int c3 = 0;
		int t3 = 0;
		for (set<string>::iterator it = mixinds.begin(); it != mixinds.end(); it++){
			//cout << *it << "\n";
			pair<int, int> p = ind2index[*it];
			t1+=2;
			if (H[p.first][i]) c1++;
			if (H[p.second][i]) c1++;
		}
		//cout << "\n";
		for (set<string>::iterator it = sourceinds.begin(); it != sourceinds.end(); it++){
			pair<int, int> p = ind2index[*it];
			t2+=2;
			if (H[p.first][i]) c2++;
			if (H[p.second][i]) c2++;
		}
		for (set<string>::iterator it = mix2.begin(); it != mix2.end(); it++){
			pair<int, int> p = ind2index[*it];
			t3+=2;
			if (H[p.first][i]) c3++;
			if (H[p.second][i]) c3++;
		}
		if (c3 == 0 or c3 == t3){
			snpinclude.push_back(false);
			snpweight.push_back(0);
			snpfreq.push_back(0);


		}
		else if (c1 == 0 and c2  == 0){
			snpinclude.push_back(false);
			snpweight.push_back(0);
			snpfreq.push_back(0);

		}
		else if (c1 == t1 and c2 == t2){
			snpinclude.push_back(false);
			snpweight.push_back(0);
			snpfreq.push_back(0);
		}
		else{
			snpinclude.push_back(true);
			double f1 = (double) c1/ (double) t1;
			double f2 = (double) c2/ (double) t2;
			double w = f1-f2;
			snpweight.push_back(w);
			snpfreq.push_back(f1);

			float tmpmix = 2* f1 * ((double) t1 - (double) c1)/ ((double) t1 - 1.0);
			float tmpsource = 2* f2 * ((double) t2 - (double) c2)/ ((double) t2 - 1.0);
			*f2_uncor  += w*w;
			*hzy_mix+= tmpmix;
			*hzy_source += tmpsource;
			*count_mix += (double) t1/ 2.0;
			*count_source += (double) t2/ 2.0;
			//cout << t1 << " "<< t2 << "\n";
			totalsnps++;
		}
	}
	return totalsnps;
}

void HapMatrix::print_all_d(set<string> inds, double maxdist){
	//cout << "here\n"; cout.flush();
	for (int i = 0; i< nsnp; i++){
		if (!snpinclude[i]) continue;
		int j = i+1;
		double gposi = genmap[i];
		while(j < nsnp and (genmap[j]-gposi) < maxdist){
			if (!snpinclude[j]){
				j++;
				continue;
			}
			double gdist = genmap[j] - gposi;
			int c11  =0;
			int c1 =0;
			int c2 =0;
			int t = 0;
			for (set<string>::iterator it = inds.begin(); it != inds.end(); it++){
				pair<int, int> p = ind2index[*it];
				t+=2;
				if (H[p.first][i] and H[p.first][j]) c11++;
				if (H[p.first][i]) c1++;
				if (H[p.first][j]) c2++;

				if (H[p.second][i] and H[p.second][j]) c11++;
				if (H[p.second][i]) c1++;
				if (H[p.second][j]) c2++;
			}
			double f11 = (double) c11 /(double)t;
			double f1 = (double) c1/ (double) t;
			double f2 = (double) c2/ (double)t;
			double d = f11-f1*f2;
			//if (c1 == t && c2 == t){
			//	j++;
			//	continue;
			//}
			//if (c1 == 0 && c2 == 0){
			//	j++;
			//	continue;
			//}
			double w = snpweight[i]* snpweight[j];
			cout << gdist << " "<< d << " "<< w<< "\n";
			j++;
		}
	}

}

float HapMatrix::calc_d(set<string> inds, int i, int j){
	float toreturn;
	int c11  =0;
	int c1 =0;
	int c2 =0;
	int t = 0;
	for (set<string>::iterator it = inds.begin(); it != inds.end(); it++){
		pair<int, int> p = ind2index[*it];
		t+=2;
		if (H[p.first][i] and H[p.first][j]) c11++;
		if (H[p.first][i]) c1++;
		if (H[p.first][j]) c2++;

		if (H[p.second][i] and H[p.second][j]) c11++;
		if (H[p.second][i]) c1++;
		if (H[p.second][j]) c2++;
	}
	double f11 = (double) c11 /(double)t;
	double f1 = (double) c1/ (double) t;
	double f2 = (double) c2/ (double)t;
	toreturn = f11-f1*f2;
	return toreturn;
}
vector<pair<vector<float>, vector<float> > > HapMatrix::bin_d(set<string> inds, double min, double max, double by){
	vector<pair<vector<float>, vector<float> > > toreturn;
	for (double i  = min; i < max; i+= by){
		pair<vector<float>, vector<float> > tmp;
		toreturn.push_back(tmp);
	}

	for (int i = 0; i< nsnp; i++){
		if (!snpinclude[i]) continue;
		int j = i+1;
		double gposi = genmap[i];
		while(j < nsnp and (genmap[j]-gposi) < max){
			if (!snpinclude[j]){
				j++;
				continue;
			}
			double gdist = genmap[j] - gposi;
			double tmpindex = gdist/by;
			float d = calc_d(inds, i, j);
			float w = snpweight[i]*snpweight[j];
			int index = (int) tmpindex;
			//cout << gdist << " "<< index << " "<< d << " "<< w << "\n";
			toreturn[index].first.push_back(d);
			toreturn[index].second.push_back(w);

			j++;
		}
	}
	return toreturn;
}

vector<vector<float> > HapMatrix::bin_d_lowmem(set<string> inds, double min, double max, double by){
	vector<vector<float> > toreturn;
	for (double i  = min; i < max; i+= by){
		vector<float> tmp (5, 0.0);
		toreturn.push_back(tmp);
	}

	for (int i = 0; i< nsnp; i++){
		if (!snpinclude[i]) continue;
		int j = i+1;
		double gposi = genmap[i];
		while(j < nsnp and (genmap[j]-gposi) < max){
			if (!snpinclude[j]){
				j++;
				continue;
			}
			double gdist = genmap[j] - gposi;
			double tmpindex = gdist/by;
			float d = calc_d(inds, i, j);
			float w = snpweight[i]*snpweight[j];
			//cout << gdist << " "<< d << " "<< w << "\n";
			int index = (int) tmpindex;
			toreturn[index].at(0)+= d;
			toreturn[index].at(1)+= w;
			toreturn[index].at(2)+= d*w;
			toreturn[index].at(3)+= w*w;
			toreturn[index].at(4)+= 1.0;


			j++;
		}
	}
	return toreturn;
}

vector<vector<float> > HapMatrix::bin_d_freq_cor_lowmem(set<string> inds, vector<pair<float, float> > cor, double min, double max, double by){
	vector<vector<float> > toreturn;
	for (double i  = min; i < max; i+= by){
		vector<float> tmp (5, 0.0);
		toreturn.push_back(tmp);
	}

	for (int i = 0; i< nsnp; i++){
		if (!snpinclude[i]) continue;
		int j = i+1;
		double gposi = genmap[i];
		while(j < nsnp and (genmap[j]-gposi) < max){
			if (!snpinclude[j]){
				j++;
				continue;
			}
			double gdist = genmap[j] - gposi;
			double tmpindex = gdist/by;
			float d = calc_d(inds, i, j);
			float w = snpweight[i]*snpweight[j];
			float fs = snpfreq[i]*snpfreq[j];
			//cout << gdist << " "<< d << " "<< w << "\n";
			int index = (int) tmpindex;
			pair<float, float> coefs = cor[index];
			double pred_d = coefs.first+ coefs.second*fs;
			d = d-pred_d;
			toreturn[index].at(0)+= d;
			toreturn[index].at(1)+= w;
			toreturn[index].at(2)+= d*w;
			toreturn[index].at(3)+= w*w;
			toreturn[index].at(4)+= 1.0;


			j++;
		}
	}
	return toreturn;
}

vector<vector<float> > HapMatrix::bin_d_freq_lowmem(set<string> inds, double min, double max, double by){
	vector<vector<float> > toreturn;
	for (double i  = min; i < max; i+= by){
		vector<float> tmp (5, 0.0);
		toreturn.push_back(tmp);
	}

	for (int i = 0; i< nsnp; i++){
		if (!snpinclude[i]) continue;
		int j = i+1;
		double gposi = genmap[i];
		while(j < nsnp and (genmap[j]-gposi) < max){
			if (!snpinclude[j]){
				j++;
				continue;
			}
			double gdist = genmap[j] - gposi;
			double tmpindex = gdist/by;
			float d = calc_d(inds, i, j);
			float w = snpfreq[i]*snpfreq[j];
			//cout << gdist << " "<< d << " "<< w << "\n";
			int index = (int) tmpindex;
			toreturn[index].at(0)+= d;
			toreturn[index].at(1)+= w;
			toreturn[index].at(2)+= d*w;
			toreturn[index].at(3)+= w*w;
			toreturn[index].at(4)+= 1.0;


			j++;
		}
	}
	return toreturn;
}
